package com.example.shardingsphere.config;

import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.*;

/**
 * @author liujie
 * @createTime 2021/8/5 14:39
 * @description
 */
@Slf4j
public class MyTableComplexAlgorithm implements ComplexKeysShardingAlgorithm<Integer> {
    @Override
    public Collection<String> doSharding(Collection<String> collection, ComplexKeysShardingValue<Integer> complexKeysShardingValue) {

        log.info("MyTableComplexAlgorithm collection: {}", collection);
        // 得到每个分片健对应的值
        Collection<Integer> orderIdValues = this.getShardingValue(complexKeysShardingValue, "order_id");

        List<String> shardingSuffix = new ArrayList<>();
        for (Integer orderId : orderIdValues) {
            String suffix = "_" + orderId % 2 ;
            for (String table : collection) {
                if (table.endsWith(suffix)) {
                    shardingSuffix.add(table);
                }
            }
        }
        return shardingSuffix;
    }

    private Collection<Integer> getShardingValue(ComplexKeysShardingValue<Integer> shardingValues, final String key) {
        Collection<Integer> valueSet = new ArrayList<>();
        Map<String, Collection<Integer>> shardingValuesMap = shardingValues.getColumnNameAndShardingValuesMap();
        Map<String, Range<Integer>> rangeValuesMap = shardingValues.getColumnNameAndRangeValuesMap();
        if (null != rangeValuesMap && rangeValuesMap.containsKey(key)) {
            Range<Integer> range = rangeValuesMap.get(key);
            // 这里取最小值和最大值
            valueSet.add(range.lowerEndpoint());
            valueSet.add(range.upperEndpoint());
        }
        if (null != shardingValuesMap && shardingValuesMap.containsKey(key)) {
            Collection<Integer> sharding = shardingValuesMap.get(key);
            valueSet.addAll(sharding);
        }
        return valueSet;
    }
}
